from io import BytesIO

import fitz
import numpy as np
from PIL import Image

from output.tools import get_data_text
from tbpu import get_parser
from tbpu.parser_tools.paragraph_parse import word_separator


def cmp(a, b):
    aLeftTopX = a['box'][0][0]
    aLeftTopY = a['box'][0][1]
    bLeftTopX = b['box'][0][0]
    bLeftTopY = b['box'][0][1]

    if aLeftTopY > bLeftTopY:
        return 1
    elif aLeftTopY < bLeftTopY:
        return -1
    else:
        if aLeftTopX > bLeftTopX:
            return 1
        elif aLeftTopX < bLeftTopX:
            return -1
        else:
            return 0

def parse_pdf_pages(pdf_path, tbpu, img_handler = None):
    page_data = parse_pdf(pdf_path, tbpu, img_handler)
    parser = get_parser(tbpu)
    pages = []
    # 计算平均置信度
    currentPage = 1
    for page in page_data:
        # 执行 tbpu
        if parser is not None and len(page) > 0:
            # page.sort(key=cmp_to_key(cmp))
            page = parser.run(page)
        pages.append({
            "content": get_data_text(page),
            "pageNum": currentPage
        })
        currentPage = currentPage + 1
    return pages

def parse_pdf(pdf_path, tbpu, img_handler = None, page_start=None, page_end=None):
    try:
        doc = fitz.open(pdf_path)
    except Exception as e:
        print(f"[Error] fitz.open error: {pdf_path}, Message: {e}")
        return []
    page_count = doc.page_count
    if page_start is not None and page_end is not None:
        if page_start < 0:
            page_start += page_count + 1
        if page_end < 0:
            page_end += page_count + 1
        if page_start < 1:
            return f"[Error] 起始页码不能小于1"
        if page_end > page_count:
            return f"[Error] 结束页码不能大于页数"
        if page_start > page_end:
            return f"[Error] 页码范围错误"
        page_list = list(range(page_start - 1, page_end))
    else:
        page_list = list(range(0, page_count))

    tb_arr = []
    for page_num in page_list:
        page = doc[page_num]
        resp = parse_single_page(page, page_num, tbpu, img_handler)

        tb_arr.append(resp)

    return tb_arr

def parse_single_page(page, page_num, tbpu, img_handler = None):
    imgs = []  # 待OCR的图片列表
    tbs = []  # text box 文本块列表
    page_rotation = page.rotation
    # 获取元素 https://pymupdf.readthedocs.io/en/latest/_images/img-textpage.png
    # 确保越界图像能被采集 https://github.com/pymupdf/PyMuPDF/issues/3171
    p = page.get_text("dict", clip=fitz.INFINITE_RECT())
    for t in p["blocks"]:  # 遍历区块（段落）
        # 图片
        if t["type"] == 1:  # 图片内容
            img_bytes = t["image"]  # 图片字节
            bbox = t["bbox"]  # 图片包围盒
            # 图片视觉大小、原始大小、缩放比例
            w1, h1 = bbox[2] - bbox[0], bbox[3] - bbox[1]
            w2, h2 = t["width"], t["height"]
            # 特殊情况：图片宽高为0
            if w2 <= 0 or h2 <= 0:
                continue
            # 单独计算宽高的缩放比例
            scale_w = w1 / w2
            scale_h = h1 / h2
            nd_arr = np.array(None)
            try:
                with Image.open(BytesIO(img_bytes)) as pimg:
                    # 记录原图格式
                    format = pimg.format
                    if not format:
                        format = "PNG"
                    if page_rotation != 0:
                        # 如果页面有旋转，逆向旋转图片字节
                        print(f"P{page_num} - 旋转 {page_rotation} °")
                        # PDF的旋转是顺时针，需要逆时针旋转图片
                        pimg = pimg.rotate(-page_rotation, expand=True)
                        # 将旋转后的图片转回bytes
                        buffered = BytesIO()
                        pimg.save(buffered, format=format)
                        img_bytes = buffered.getvalue()
                    nd_arr = np.array(pimg)
            except Exception:
                print("转换文档图片异常。")
            # 记录图片
            if img_handler is not None:
                imgs.append(
                    {
                        "nd_arr": nd_arr,
                        "bytes": img_bytes,
                        "xy": (bbox[0], bbox[1]),
                        "scale_w": scale_w,
                        "scale_h": scale_h,
                        "end": "\n"
                    }
                )
        # 文本
        elif t["type"] == 0:
            l = len(t["lines"]) - 1
            for index, line in enumerate(t["lines"]):  # 遍历每一行
                text = ""
                for span in line["spans"]:  # 遍历每一文本块
                    text += span["text"]
                if text:
                    b = line["bbox"] #(左上x 左上y 右下x 右下y)
                    tb = {
                        "box": [
                            [b[0], b[1]],
                            [b[2], b[1]],
                            [b[2], b[3]],
                            [b[0], b[3]],
                        ],
                        "text": text,
                        "score": 1,
                        "end": "\n" if index == l else "",  # 结尾符
                        "from": "text",  # 来源：直接提取文本
                    }
                    tbs.append(tb)
    # 补充结尾符
    for i1 in range(len(tbs) - 1):
        if tbs[i1]["end"]:  # 跳过已有结尾符的
            continue
        i2 = i1 + 1
        sep = word_separator(tbs[i1]["text"][-1], tbs[i2]["text"][0])
        tbs[i1]["end"] = sep

    if imgs:
        for img in imgs:
            res = img_handler(img['nd_arr'])
            if res is not None:
                x, y = img["xy"]
                scale_w = img["scale_w"]
                scale_h = img["scale_h"]
                for r in res:
                    # 将所有文本块的坐标，从图片相对坐标系，转为页面绝对坐标系
                    for bi in range(4):
                        r["box"][bi][0] = r["box"][bi][0] * scale_w + x
                        r["box"][bi][1] = r["box"][bi][1] * scale_h + y
                    r["from"] = "img"  # 来源：OCR
                    tbs.append(r)
    parser = get_parser(tbpu)
    if parser is not None and len(tbs) > 0:
        tbs = parser.run(tbs)
    return tbs
